#include "RLE.h"

//
// constants
//
#define MAX_BLOCK_SIZE       127
#define MAX_COMP_BLOCK_SIZE  130

#define COMP_BLOCK_PREFIX    (1 << 7)
#define INDICATOR_MASK       0x7F

//
// Purpose:
//   Compresses the specified buffer and pastes it into the
//   input buffer if the compression did benefit. Otherwise
//   FALSE is returned and the input buffer remains unchanged.
//
BOOL RLE::Encode( PBYTE pData, DWORD cb, DWORD &cbNew )
{
	UINT    iOut        = 4;
	UINT    iCur        = 0;
	UINT    u, i;
	DWORD   cbNewBlock;
	BOOL    bBenefits   = TRUE;
	PBYTE   pOut        = (PBYTE)malloc( cb );
	PDWORD  pdwOrgSize  = (PDWORD)pOut; // original size is pasted at the top of the block

	if ( !pOut )
		return FALSE; // ERR

	//
	// enter main loop
	//
	while ( TRUE )
	{
		//
		// search the start of the next compressible block
		//
		for ( u=iCur; u < cb; u++ )
		{
			// how many equal symbols after each other at the current location ?
			i          = u;
			while ( i+1 < cb && pData[i] == pData[i+1])
				++i;
			if ( (i - u + 1) >= 3 ) // at least 3 equal bytes in a row ?
				break;
		}

		//
		// stamp block(s) for the uncompressible bytes
		//
		cbNewBlock = u - iCur;

		// would action exceed input buffers range ?
		if ( iOut + cbNewBlock >= cb )
		{
			bBenefits = FALSE; // ERR
			break;
		}

		while ( cbNewBlock )
		{
			DWORD cbThisBlock = min( cbNewBlock, MAX_BLOCK_SIZE );
			pOut[iOut++]      = (BYTE)(cbThisBlock);
			memcpy(
				&pOut[iOut],
				&pData[iCur],
				cbThisBlock );

			iOut             += cbThisBlock;
			iCur             += cbThisBlock;
			cbNewBlock       -= cbThisBlock;
		}

		//
		// now print the compressible block
		//
		cbNewBlock = i - u + 1;

		// would action exceed input buffers range ?
		if ( iOut + 2 >= cb )
		{
			bBenefits = FALSE; // ERR
			break;
		}

		while ( cbNewBlock >= 3 )
		{
			DWORD cbThisBlock = min( cbNewBlock, MAX_COMP_BLOCK_SIZE );
			pOut[iOut++]      = (BYTE)( (cbThisBlock - 3) | COMP_BLOCK_PREFIX );
			pOut[iOut++]      = pData[u];

			iCur             += cbThisBlock;
			cbNewBlock       -= cbThisBlock;
		}

		// end of input reached ?
		if ( iCur == cb )
			break;
	}

	// post-initialize output variable
	if ( bBenefits )
	{
		// paste size to compressed block
		*pdwOrgSize  = cb;
		// size of compressed data block == current index in it
		cbNew        = iOut;
		// overwrite input buffer with the compressed block
		memcpy(
			pData,
			pOut,
			cbNew );
	}

	//
	// tidy up
	//
	free( pOut );

	// return result
	return bBenefits;
}

//
// Purpose:
//   Decompress the specified buffer, encoded with RLE::Encode,
//   and pastes the result into the input buffer.
//
// Returns:
//   Only negative if there wasn't enough memory available
//
BOOL RLE::Decode( PBYTE & pData, DWORD & cb )
{
	DWORD  cbTotal  = *(PDWORD)pData;
	PBYTE  pTmp     = (PBYTE)malloc( cbTotal );
	UINT   iCur     = 4;
	UINT   iOut     = 0;

	if ( !pTmp )
		return FALSE; // ERR

	//
	// enter main loop
	//
	while ( iOut != cbTotal )
	{
		BYTE byIndicator = pData[iCur++];
		BYTE byCount     = byIndicator & INDICATOR_MASK;

		if ( byIndicator & COMP_BLOCK_PREFIX )
		{
			//
			// decode compressed block
			//
			byCount += 3;
			memset(
				&pTmp[iOut],
				pData[iCur++],
				byCount );
			iOut += byCount;
		}
		else
		{
			//
			// decode UN-compressed block
			//
			memcpy(
				&pTmp[iOut],
				&pData[iCur],
				byCount );
			iOut += byCount;
			iCur += byCount;
		}
	}

	//
	// apply decompressed data to input buffer
	//
	pData = (PBYTE)realloc( pData, cbTotal );
	if ( !pData )
	{
		free( pTmp );
		return FALSE; // ERR
	}
	memcpy( pData, pTmp, cbTotal );

	//
	// tidy up
	//
	free( pTmp );

	return TRUE; // OK
}